import cv2
import os
import numpy as np
from ultralytics import YOLO

# 配置参数
input_folder = "D:/Desk/test1/content/"  # 原始图片目录
output_folder = "D:/Desk/test1/result/"  # 输出目录
target_ratio = 3/4  # 宽高比 3:4
model_path = "../yolov8x-cls.pt"  # 或专用服装检测模型

# 加载预训练模型（需先下载）
model = YOLO(model_path)  # 使用YOLOv8分类模型（或换成服装检测模型）

def detect_clothing_center(img):
    """使用YOLO模型定位衣物中心"""
    results = model(img)  # 推理
    
    # 解析结果：寻找概率最高的衣物区域
    max_conf = 0
    best_box = None
    
    for result in results:
        if result.boxes is not None:
            for box in result.boxes:
                cls_id = int(box.cls)
                conf = box.conf.item()
                # 筛选衣物类别（根据模型类别ID调整）
                if cls_id in [16, 18, 20]:  # 16=人，18=包，20=上衣（COCO类别示例）
                    if conf > max_conf:
                        max_conf = conf
                        best_box = box.xyxy[0].cpu().numpy()
    
    if best_box is not None:
        x1, y1, x2, y2 = best_box.astype(int)
        return ((x1 + x2) // 2, (y1 + y2) // 2)
    else:
        # 回退策略：使用人体检测
        hog = cv2.HOGDescriptor()
        hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
        regions, _ = hog.detectMultiScale(img)
        if len(regions) > 0:
            x, y, w, h = max(regions, key=lambda r: r[2]*r[3])
            return (x + w//2, y + h//2)
        else:
            return (img.shape[1]//2, img.shape[0]//2)  # 默认居中

# 批量处理
os.makedirs(output_folder, exist_ok=True)
for filename in os.listdir(input_folder):
    if not filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        continue

    img_path = os.path.join(input_folder, filename)
    img = cv2.imread(img_path)
    if img is None:
        continue

    h, w = img.shape[:2]
    target_width = int(target_ratio * h)
    
    # 获取中心点
    center_x, _ = detect_clothing_center(img)
    
    # 动态裁剪
    x_start = max(0, center_x - target_width//2)
    x_end = x_start + target_width
    
    if x_end > w:
        x_start = max(0, w - target_width)
        x_end = w

    cropped = img[0:h, x_start:x_end]
    cv2.imwrite(os.path.join(output_folder, filename), cropped)

print("处理完成！")